"""
MCP上下文管理器

负责ModelContext的生命周期管理和持久化
"""

from typing import Dict, Any, Optional, List
from datetime import datetime, timedelta
from sqlalchemy.orm import Session
from app.models.mcp_context import MCPContext
from app.models.mcp_session import MCPSession
from app.core.logging_config import app_logger
import json


class MCPContextManager:
    """
    MCP上下文管理器
    
    功能：
    1. 创建、更新、查询、删除上下文
    2. 上下文版本管理
    3. 上下文持久化到MySQL
    4. 上下文缓存到Redis（可选）
    """
    
    def __init__(self, db: Session, redis_client: Optional[Any] = None):
        """
        初始化上下文管理器
        
        Args:
            db: 数据库会话
            redis_client: Redis客户端（可选）
        """
        self.db = db
        self.redis = redis_client
        self.logger = app_logger
    
    def create_context(
        self,
        session_id: str,
        context_data: Dict[str, Any],
        context_type: str = "user",
        context_metadata: Optional[Dict[str, Any]] = None,
        model_id: Optional[str] = None
    ) -> MCPContext:
        """
        创建新的上下文
        
        Args:
            session_id: MCP会话ID
            context_data: 上下文数据
            context_type: 上下文类型（system/user/assistant/tool）
            context_metadata: 元数据
            model_id: 模型ID
            
        Returns:
            MCPContext: 创建的上下文对象
        """
        # 计算数据大小
        context_json = json.dumps(context_data)
        size_bytes = len(context_json.encode('utf-8'))
        
        # 创建上下文记录
        context = MCPContext(
            session_id=session_id,
            context_data=context_data,
            context_type=context_type,
            context_metadata=context_metadata or {},
            model_id=model_id,
            version=1,
            size_bytes=size_bytes,
            is_encrypted=False
        )
        
        self.db.add(context)
        self.db.commit()
        self.db.refresh(context)
        
        self.logger.info(f"创建MCP上下文: session_id={session_id}, id={context.id}")
        
        return context
    
    def update_context(
        self,
        context_id: int,
        context_data: Dict[str, Any],
        increment_version: bool = True
    ) -> MCPContext:
        """
        更新上下文
        
        Args:
            context_id: 上下文ID
            context_data: 新的上下文数据
            increment_version: 是否增加版本号
            
        Returns:
            MCPContext: 更新后的上下文对象
        """
        context = self.db.query(MCPContext).filter(MCPContext.id == context_id).first()
        
        if not context:
            raise ValueError(f"上下文不存在: id={context_id}")
        
        # 更新数据
        context.context_data = context_data
        context_json = json.dumps(context_data)
        context.size_bytes = len(context_json.encode('utf-8'))
        
        if increment_version:
            context.version += 1
        
        self.db.commit()
        self.db.refresh(context)
        
        self.logger.info(f"更新MCP上下文: id={context_id}, version={context.version}")
        
        return context
    
    def get_context(self, context_id: int) -> Optional[MCPContext]:
        """获取上下文"""
        context = self.db.query(MCPContext).filter(MCPContext.id == context_id).first()
        return context
    
    def get_session_contexts(
        self,
        session_id: str,
        context_type: Optional[str] = None,
        limit: int = 50
    ) -> List[MCPContext]:
        """
        获取会话的所有上下文
        
        Args:
            session_id: 会话ID
            context_type: 上下文类型过滤（可选）
            limit: 返回数量限制
            
        Returns:
            List[MCPContext]: 上下文列表
        """
        query = self.db.query(MCPContext).filter(MCPContext.session_id == session_id)
        
        if context_type:
            query = query.filter(MCPContext.context_type == context_type)
        
        contexts = query.order_by(MCPContext.created_at.desc()).limit(limit).all()
        
        return contexts
    
    def delete_context(self, context_id: int) -> bool:
        """删除上下文"""
        context = self.db.query(MCPContext).filter(MCPContext.id == context_id).first()
        
        if not context:
            return False
        
        session_id = context.session_id
        
        self.db.delete(context)
        self.db.commit()
        
        self.logger.info(f"删除MCP上下文: id={context_id}, session_id={session_id}")
        
        return True
    
    def delete_session_contexts(self, session_id: str) -> int:
        """
        删除会话的所有上下文
        
        Args:
            session_id: 会话ID
            
        Returns:
            int: 删除的上下文数量
        """
        deleted_count = self.db.query(MCPContext).filter(
            MCPContext.session_id == session_id
        ).delete()
        
        self.db.commit()
        
        self.logger.info(f"删除会话所有上下文: session_id={session_id}, count={deleted_count}")
        
        return deleted_count

